#include <memory.h>
#include <stdio.h>
#include <winsock2.h>  // This must be included before "windows.h"
#include <windows.h>

#include "checksum.h"

// 
// Update checksum according to old and new port number,
// all parameters (including the checksum) must be in host
// byte order, while the new checksum returned is in network 
// byte order and do not need to be re-ordered.
//
u_short checksum_adjust( u_short chksum, u_short oldp, u_short newp )
{
    long x;
    u_short ret;

    x = chksum;
    x = ~x & 0xffff;
    x -= oldp & 0xffff;
    if( x <= 0 ) {
	    x--;
	    x &= 0xffff;
    }

    x += newp & 0xffff;
    if( x & 0x10000 ) {
	    x++;
	    x &= 0xffff;
    }

    x = ~x & 0xffff;
    ret = (u_short)( x/256 );
    ret += (u_short)( (x & 0xff) << 8 );
    return ret;
}

//
// Checksum computing function for raw data without folding
// 32bits(u_long) result into 16bits(u_short) checksum
//
u_long checksum_unfold( u_short *buffer, int size )
{
    u_long x = 0;
    
    while( size > 1 )
    {
        x += *buffer++;
        size -= sizeof(u_short);
    }
    
    if (size > 0)
    {
        x += ( *(u_char *)buffer & 0xff ) ;
    }
    
    return x;
}

//
// TCP checksum for IPv4 packet
//
void checksum_tcp4( ip_header *ih, tcp_header *th )
{
    u_short checksum;
	// u_short checksum0;
    u_long  w1, w2, w3;
    u_short size;
    psd_header psdh;
    
    memset( &psdh, 0, sizeof(psd_header)  );
    memcpy( psdh.saddr, ih->saddr, 4 );
	memcpy( psdh.daddr, ih->daddr, 4 );
    
    psdh.protocol = ih->protocol;
    size = ntohs(ih->length) - (ih->ver_ihl & 0x0f) * 4;
    psdh.length = htons(size);
    
    checksum = th->checksum;
	printf("checksum_tcp4: Old TCP checksum: %02x\n", checksum);
    th->checksum = 0;
    
    // Compute psuedo-header and tcp data seperately
    w1 = checksum_unfold( (u_short *)&psdh, sizeof(psd_header) );
    w2 = checksum_unfold( (u_short *)th, size );
    w3 = w1 + w2;
    
    // Fold the final result into 16 bits
    while(w3>>16)
        w3 = (w3>>16) + (w3 & 0xffff);
    checksum = (u_short)(~w3);
    
    printf("checksum_tcp4: Recompute TCP checksum: %02x\n", checksum);
    th->checksum = checksum;
    
    // Compute ip header checksum
    checksum = ih->checksum;
    printf("checksum_tcp4: Old IP checksum: %02x\n", checksum);
    size = (ih->ver_ihl & 0x0f) * 4;
    
    ih->checksum = 0;
    w1 = checksum_unfold( (u_short *)ih, size );
    while (w1>>16)
        w1 = (w1>>16) + (w1 & 0xffff);
    checksum = (u_short)(~w1);
    printf("checksum_tcp4: Recompute IP checksum: %02x\n", checksum);
    ih->checksum = checksum;
    
    return;
}

//
// UDP checksum for IPv4 packet
//
void checksum_udp4( ip_header *ih, udp_header *uh )
{
    u_short checksum;
	//u_short checksum0;
    u_long  w1, w2, w3;
    u_short size;
    psd_header psdh;
    
    memset( &psdh, 0, sizeof(psd_header) );
    memcpy( psdh.saddr, ih->saddr, 4 );
	memcpy( psdh.daddr, ih->daddr, 4 );
    
    psdh.protocol = ih->protocol;
    psdh.length = uh->length;
    
    size = ntohs(uh->length);
    checksum = uh->checksum;
    printf("checksum_udp4: Old UDP checksum: %02x\n", checksum);
    uh->checksum = 0;
    
    // Compute psuedo-header and udp data seperately
    w1 = checksum_unfold( (u_short *)&psdh, sizeof(psd_header) );
    w2 = checksum_unfold( (u_short *)uh, size );
    w3 = w1 + w2;
    
    // Fold the final result into 16 bits
    while(w3>>16)
        w3 = (w3>>16) + (w3 & 0xffff);
    checksum = (u_short)(~w3);
    
    printf("checksum_udp4: Recompute checksum: %02x\n", checksum);
    uh->checksum = checksum;
    
    // Compute ip header checksum
    checksum = ih->checksum;
    printf("checksum_udp4: Old IP checksum: %02x\n", checksum);
    size = (ih->ver_ihl & 0x0f) * 4;
    
    ih->checksum = 0;
    w1 = checksum_unfold( (u_short *)ih, size );
    while (w1>>16)
        w1 = (w1>>16) + (w1 & 0xffff);
    checksum = (u_short)(~w1);
    printf("checksum_udp4: Recompute IP checksum: %02x\n", checksum);
    ih->checksum = checksum;
    
    return;
}

//
// ICMP checksum for IPv4 packet
//
void checksum_icmp4( ip_header *ih, icmp_header *icmph )
{
    // u_short checksum;
	// u_short	checksum0;
    u_long  w1;
    u_short size;

	// Compute icmp checksum

    //checksum = icmph->checksum;
	//printf("checksum_icmp4: Old ICMP checksum: %02x\n", checksum);
	size = ntohs(ih->length) - (ih->ver_ihl & 0x0F) * 4;
    
	icmph->checksum = 0;
    w1 = checksum_unfold( (u_short *)icmph, size );
    while(w1>>16)
        w1 = (w1>>16) + (w1 & 0xffff);
    //checksum = (u_short)(~w1);
    //printf("checksum_icmp4: Recompute ICMP checksum: %02x\n", checksum);
    //icmph->checksum = checksum;
	icmph->checksum = (u_short)(~w1);
    
    // Compute ip header checksum

    //checksum = ih->checksum;
    //printf("checksum_icmp4: Old IP checksum: %02x\n", checksum);
    size = (ih->ver_ihl & 0x0F) * 4;
    
    ih->checksum = 0;
    w1 = checksum_unfold( (u_short *)ih, size );
    while (w1>>16)
        w1 = (w1>>16) + (w1 & 0xffff);
    //checksum = (u_short)(~w1);
    //printf("checksum_icmp4: Recompute IP checksum: %02x\n", checksum);
    //ih->checksum = checksum;
    ih->checksum = (u_short)(~w1);

    return;
}

//
// TCP checksum for IPv6 packet
//
void checksum_tcp6( ip6_header *ip6h, tcp_header *th )
{
    u_short checksum;
	// u_short checksum0;
    u_long  w1, w2, w3;
    u_short size;
    psd6_header psd6h;
    
    memset( &psd6h, 0, sizeof(psd6_header) );
    memcpy( psd6h.saddr, ip6h->saddr, 16 );
    memcpy( psd6h.daddr, ip6h->daddr, 16 );
    
    psd6h.length[1] = ip6h->payload;
    psd6h.nexthdr = ip6h->nexthdr;
    
    size = ntohs(ip6h->payload);
    checksum = th->checksum;
	printf("checksum_tcp6: Old TCP checksum: %02x\n", checksum);
    th->checksum = 0;
    
    // Compute psuedo-header and tcp data seperately
    w1 = checksum_unfold( (u_short *)&psd6h, sizeof(psd6_header) );
    w2 = checksum_unfold( (u_short *)th, size );
    w3 = w1 + w2;
    
    // Fold the final result into 16 bits
    while(w3>>16)
        w3 = (w3>>16) + (w3 & 0xffff);
    checksum = (u_short)(~w3);
    
    printf("checksum_tcp6: Recompute TCP checksum: %02x\n", checksum);
    th->checksum = checksum;
    
    return;
}

//
// UDP checksum for IPv6 packet
//
void checksum_udp6( ip6_header *ip6h, udp_header *uh )
{
    u_short checksum;
	// u_short checksum0;
    u_long  w1, w2, w3;
    u_short size;
    psd6_header psd6h;
    
    memset( &psd6h, 0, sizeof(psd6_header) );
    memcpy( psd6h.saddr, ip6h->saddr, 16 );
    memcpy( psd6h.daddr, ip6h->daddr, 16 );
    
    psd6h.length[1] = ip6h->payload;
    psd6h.nexthdr = ip6h->nexthdr;
    
    size = ntohs(ip6h->payload);
    checksum = uh->checksum;
	printf("checksum_udp6: Old UDP checksum: %02x\n", checksum);
    uh->checksum = 0;
    
    // Compute psuedo-header and udp data seperately
    w1 = checksum_unfold( (u_short *)&psd6h, sizeof(psd6_header) );
    w2 = checksum_unfold( (u_short *)uh, size );
    w3 = w1 + w2;
    
    // Fold the final result into 16 bits
    while(w3>>16)
        w3 = (w3>>16) + (w3 & 0xffff);
    checksum = (u_short)(~w3);
    
    printf("checksum_udp6: Recompute UDP checksum: %02x\n", checksum);
    uh->checksum = checksum;
    
    return;
}

//
// ICMPv6 checksum for IPv6 packet
//
void checksum_icmp6( ip6_header *ip6h, icmp6_header *icmp6h )
{
    // u_short checksum;
	// u_short checksum0;
    u_long  w1, w2, w3;
    u_short size;
    psd6_header psd6h;
    
    memset( &psd6h, 0, sizeof(psd6_header) );
    memcpy( psd6h.saddr, ip6h->saddr, 16 );
    memcpy( psd6h.daddr, ip6h->daddr, 16 );
    
    psd6h.length[1] = ip6h->payload;
    psd6h.nexthdr = ip6h->nexthdr;
    
    size = ntohs(ip6h->payload);
    //checksum = icmp6h->checksum;
	//printf("checksum_icmp6: Old ICMPv6 checksum: %02x\n", checksum);
    icmp6h->checksum = 0;
    
    // Compute psuedo-header and icmpv6 data seperately
    w1 = checksum_unfold( (u_short *)&psd6h, sizeof(psd6_header) );
    w2 = checksum_unfold( (u_short *)icmp6h, size );
    w3 = w1 + w2;
    
    // Fold the final result into 16 bits
    while(w3>>16)
        w3 = (w3>>16) + (w3 & 0xFFFF);
    //checksum = (u_short)(~w3);
    //printf("checksum_icmp6: Recompute ICMPv6 checksum: %02x\n", checksum);
    //icmp6h->checksum = checksum;
    icmp6h->checksum = (u_short)(~w3);

    return;
}
